package org.walklink.agent.track;

import com.alibaba.ttl.TransmittableThreadLocal;

import java.util.Stack;

/**
 * 追踪管理器
 *
 * @author lidishan
 */
public class TrackManager {

    private static final TransmittableThreadLocal<Stack<TraceInfo>> TRACK = new TransmittableThreadLocal<>();

    private static TraceInfo createSpan() {
        Stack<TraceInfo> stack = TRACK.get();
        if (stack == null) {
            stack = new Stack<>();
            TRACK.set(stack);
        }
        String linkId;
        if (stack.isEmpty()) {
            linkId = TrackContext.getTraceInfo().getTraceId();
            if (linkId == null) {
                linkId = "nvl";
                TrackContext.setTraceInfo(new TraceInfo(linkId));
            }
        } else {
            TraceInfo span = stack.peek();
            linkId = span.getTraceId();
            TrackContext.setTraceInfo(new TraceInfo(linkId));
        }
        return new TraceInfo(linkId);
    }

    public static void createEntrySpan() {
        TraceInfo span = createSpan();
        Stack<TraceInfo> stack = TRACK.get();
        stack.push(span);
    }


    public static TraceInfo getExitSpan() {
        Stack<TraceInfo> stack = TRACK.get();
        if (stack == null || stack.isEmpty()) {
            TrackContext.clear();
            return null;
        }
        return stack.pop();
    }

    public static TraceInfo getCurrentSpan() {
        Stack<TraceInfo> stack = TRACK.get();
        if (stack == null || stack.isEmpty()) {
            return null;
        }
        return stack.peek();
    }


}